// SPDX-License-Identifier: GPL-2.0

#include <linux/module.h>
#include <linux/rbtree.h>
#include <linux/slab.h>
#include "dim_rb_tree.h"

struct dim_rb_node {
	struct rb_node node_st;
	void *data;
};

struct dim_rb_root {
	struct rb_root root;
	struct rb_node *cur;
	dim_rb_compare comp_func;
};

dim_rb_handle dim_rb_create(dim_rb_compare comp_func)
{
	struct dim_rb_root *root = NULL;

	if (comp_func == NULL)
		return NULL;

	root = kmalloc(sizeof(*root), GFP_KERNEL);
	if (root == NULL)
		return NULL;

	root->root.rb_node = NULL;
	root->cur = NULL;
	root->comp_func = comp_func;

	return (dim_rb_handle)root;
}

void dim_rb_destroy(dim_rb_handle root)
{
	if (root != NULL)
		kfree(root);
}

static struct dim_rb_node *rb_find_node(dim_rb_handle root, void *data)
{
	struct rb_node *node = NULL;
	struct dim_rb_root *dim_root = (struct dim_rb_root *)root;
	struct dim_rb_node *find = NULL;
	int result;

	if (dim_root == NULL)
		return NULL;

	node = dim_root->root.rb_node;
	while (node != NULL) {
		find = rb_entry(node, struct dim_rb_node, node_st);
		if (find == NULL)
			return NULL;

		result = dim_root->comp_func(data, find->data);
		if (result < 0)
			node = node->rb_left;
		else if (result > 0)
			node = node->rb_right;
		else
			return find;
	}

	return NULL;
}

void *dim_rb_find_data(dim_rb_handle root, void *data)
{
	struct dim_rb_node *node = rb_find_node(root, data);

	if (node == NULL)
		return NULL;

	return node->data;
}

int dim_rb_insert_data(dim_rb_handle root, void *new_data, void **find_data)
{
	struct dim_rb_root *dim_root = (struct dim_rb_root *)root;
	struct rb_node **new = NULL;
	struct rb_node *parent = NULL;
	struct dim_rb_node *node = NULL;
	struct dim_rb_node *new_node = NULL;
	int result;

	if (dim_root == NULL || new_data == NULL)
		return -EINVAL;

	new = &(dim_root->root.rb_node);
	while (*new != NULL) {
		node = rb_entry(*new, struct dim_rb_node, node_st);
		if (node == NULL)
			return -EFAULT;
		result = dim_root->comp_func(new_data, node->data);

		parent = *new;
		if (result < 0) {
			new = &((*new)->rb_left);
		} else if (result > 0) {
			new = &((*new)->rb_right);
		} else {
			if (find_data != NULL)
				*find_data = node->data;
			return -EEXIST;
		}
	}

	new_node = kzalloc(sizeof(*new_node), GFP_KERNEL);
	if (new_node == NULL)
		return -ENOMEM;

	new_node->data = new_data;
	rb_link_node(&new_node->node_st, parent, new);
	rb_insert_color(&new_node->node_st, root);

	return 0;
}

void *dim_rb_first_data(dim_rb_handle root)
{
	struct dim_rb_root *dim_root = (struct dim_rb_root *)root;
	struct dim_rb_node *find = NULL;

	if (dim_root == NULL)
		return NULL;

	dim_root->cur = rb_first(&dim_root->root);

	if (dim_root->cur == NULL)
		return NULL;

	find = rb_entry(dim_root->cur, struct dim_rb_node, node_st);

	return find == NULL ? NULL : find->data;
}

void *dim_rb_next_data(dim_rb_handle root)
{
	struct dim_rb_root *dim_root = (struct dim_rb_root *)root;
	struct dim_rb_node *find = NULL;

	if (dim_root == NULL || dim_root->cur == NULL)
		return NULL;

	dim_root->cur = rb_next(dim_root->cur);
	if (dim_root->cur == NULL)
		return NULL;

	find = rb_entry(dim_root->cur, struct dim_rb_node, node_st);

	return find == NULL ? NULL : find->data;
}

void *dim_rb_remove_first_data(dim_rb_handle root)
{
	struct dim_rb_root *dim_root = (struct dim_rb_root *)root;
	struct dim_rb_node *find = NULL;
	struct rb_node *node = NULL;
	void *data = NULL;

	if (dim_root == NULL)
		return NULL;

	node = rb_first(&dim_root->root);
	if (node == NULL)
		return NULL;

	find = rb_entry(node, struct dim_rb_node, node_st);
	rb_erase(node, &dim_root->root);
	data = find == NULL ? NULL : find->data;

	if (find != NULL)
		kfree(find);

	return data;
}
